   
# ---------------------------------------------------------------
## Scripts to load CITUP and PyClone to plot trees (Figure 2)
## Author: Veronique LeBlanc
## https://github.com/vleblanc/GBM-PDO-paper/blob/main/WES/fig2_trees.R
## https://github.com/hanhanial/PyClone_VI/blob/c2b7af0eb01d2f1e5fcadf907948bec3b5cb98d1/scripts/S05_plotting_citup-clone-trees.r
# ---------------------------------------------------------------

#library(ggplot2)
library(argparser)
library(reshape2)
library(data.table)
library(dplyr)


##############################################################################

argp <- arg_parser("Plot the CitupPlot")
argp <- add_argument(argp, "--work_dir" , help="")
argp <- add_argument(argp, "--pyclone_path" , help="")
argp <- add_argument(argp, "--smg_gene_file" , help="")
argp <- add_argument(argp, "--output_path" , help="")
argp <- add_argument(argp, "--sample_file" , help="")
argp <- add_argument(argp, "--sample_msi_file" , help="")
argp <- add_argument(argp, "--class_order" , help="")
argp <- add_argument(argp, "--class_a_order" , help="")
argp <- add_argument(argp, "--clone_t" , help="")
argp <- add_argument(argp, "--ccf_file" , help="")

argv <- parse_args(argp)

work_dir <- argv$work_dir
pyclone_path <- argv$pyclone_path
smg_gene_file <- argv$smg_gene_file
output_path <- argv$output_path
class_order <- argv$class_order
class_a_order <- argv$class_a_order
sample_file <- argv$sample_file
sample_msi_file <- argv$sample_msi_file
clone_t <- as.numeric(argv$clone_t)
ccf_file <- argv$ccf_file

if(1!=1){

  work_dir <- "~/20220915_gastric_multiple/dna_combinePublic"
  pyclone_path <- "~/20220915_gastric_multiple/dna_combinePublic/Pyclone/CloneImages"
  smg_gene_file <- paste0( work_dir , "/public_ref/importTantGene.addCGC.list")
	output_path <- "~/20220915_gastric_multiple/dna_combinePublic/Pyclone/selectGCClone"
  ccf_file <- paste(work_dir,"/mutationTime/result/All_CCF_mutTime.tsv",sep="")
  clone_t <- 0.6
  class_order <- paste(work_dir,"/config/Class_order.list",sep="")
  class_a_order <- paste(work_dir,"/config/Class_order_sub.list",sep="")
  sample_file <- paste(work_dir,"/config/tumor_normal.class.list",sep="")
  sample_msi_file <- paste(work_dir,"/config/tumor_normal.class.MSI.list",sep="")

}

dir.create(output_path , recursive = T)

##############################################################################
# important gene
#dat_gene <- data.frame(fread(gene_list))
dat_smg <- data.frame(fread(smg_gene_file))
#geneShow_list <- unique(c(dat_smg$Gene_Symbol , dat_gene$gene))

##############################################################################

class_order <- data.frame(fread(class_order))
class_a_order <- data.frame(fread(class_a_order))
info <- data.frame(fread(sample_file))
info_msi <- data.frame(fread(sample_msi_file))
dat_ccf <- data.frame(fread(ccf_file))

## 5个样本也纳入
#info <- subset( info , Type != "IM + IGC + DGC" )

##############################################################################

info$ms_type <- "MSS"
#info_msi$ms_type <- "MSI"
#info <- rbind( info , info_msi )

##############################################################################

Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

##############################################################################
## 取每个样本，在胃癌为克隆的突变
## cluster均是经过质控
result <- c()
result_all <- c()

for(Normal in unique(info$Normal)){

  print(Normal)

  # loci
  pyclone_loci_qc_file <- paste0(pyclone_path , "/" , Normal , "_Pyclone_cluster.tsv")
  if(  file.exists(pyclone_loci_qc_file) ){
    # pyclone loci
    dat_loci <- read.delim(pyclone_loci_qc_file)
    dat_loci$Variant_Classification <- sapply( strsplit(dat_loci$mutation_id , ":") , "[" , 2 )
    dat_loci$gene <- sapply( strsplit(dat_loci$mutation_id , ":") , "[" , 1 )
    # 功能性突变
    dat_loci <- subset(dat_loci , Variant_Classification %in% Variant_Type )
    if(nrow(dat_loci) > 0 ){
      dat_loci$Normal <- Normal

      class_type <- unique(dat_loci$Class)[unique(dat_loci$Class)!="IM"]
      if( length(unique(dat_loci$Class)[unique(dat_loci$Class)!="IM"]) > 1 ){
        dat_loci$Type <- "IGC_DGC"
      }else{
        dat_loci$Type <- class_type
      }

      result_all <- rbind(dat_loci , result_all)

      # 突变在胃癌中为主克隆
      dat_metClone <- subset( dat_loci , cellular_prevalence >= clone_t & Class %in% c("IGC" , "DGC") )
      if(nrow(dat_metClone) > 0){
        result <- rbind(result , dat_metClone)
      }
    }
  }
}

## 统计突变的人数
result2 <- result %>%
group_by(gene) %>%
summarize(CloneMutNum=length(unique(Normal)) , Normal = paste0(unique(Normal) , collapse=",") , Class = paste0(unique(Class) , collapse="_") )
result2$Class <- ifelse( result2$Class == "DGC_IGC" , "IGC_DGC" , result2$Class )

## 突变至少在2个样本中（用于非报道的基因筛选）
result3 <- result2[result2$CloneMutNum >= 2 , ]
result4 <- subset( result2 , gene %in% dat_smg$Gene_Symbol & CloneMutNum >= 1 )
result3 <- unique(rbind( result3 , result4 ))

result_record <- result2
result_record$Filter1 <- ifelse( result_record$gene %in% dat_smg$Gene_Symbol , "SMG" , "" )
result_record$Filter1 <- ifelse( result_record$CloneMutNum >= 2 & !(result_record$gene %in% dat_smg$Gene_Symbol) , "NoSMG_2Clone" , result_record$Filter1 )

##############################################################################
## 检查该基因所在的克隆簇是否与其它经典的driver在一起
result_gene <- c()
for(geneN in result_record$gene){

  tmp <- subset(result , gene == geneN)
  for(Normal in unique(tmp$Norma)){
    cluster_id <- unique(tmp[tmp$Normal==Normal,"cluster_id"])

    ## 已报道的driver基因是否在同一簇
    tmp_genes <- unique(result_all[which(result_all$cluster_id==cluster_id & result_all$Normal==Normal),"gene"])
    tmp_genes <- paste0(tmp_genes[tmp_genes %in% dat_smg$Gene_Symbol] , collapse = ",")

    tmp_res <- data.frame( gene = geneN , Normal = Normal , cluster_id = cluster_id , clusterSMG = tmp_genes)
    result_gene <- rbind(result_gene , tmp_res)
  }
}

images_name <- paste0(output_path , "/GCClone_gene.AllWithDriver.tsv" )
write.table(result_gene , images_name , sep = '\t' , quote=F , row.names=F)

##############################################################################
## 筛选标准
## 1、其为经典的driver
## 2、若非经典的driver:存在克隆簇不和经典的driver基因一起出现
gene_list1 <- unique(subset(result_gene , clusterSMG=="")$gene)
gene_list2 <- unique(result_gene[result_gene$gene %in% dat_smg$Gene_Symbol,"gene"])
gene_list4 <- unique(c(gene_list1 , gene_list2))

## 克隆演化基因的集合
result_gene2 <- result_gene[result_gene$gene %in% gene_list4,]
result_gene2$SMG <- ifelse( result_gene2$gene %in% dat_smg$Gene_Symbol , "SMG" , "" )

## 输出
images_name <- paste0(output_path , "/GCClone_gene.tsv" )
write.table(result_gene2 , images_name , sep = '\t' , quote=F , row.names=F)

##############################################################################

result_record$Filter2 <- ifelse( result_record$gene %in% gene_list1 & result_record$CloneMutNum >= 2 , "Clone_Independent & NoSMG_2Clone" , "" )
result_record$Filter2 <- ifelse( result_record$gene %in% gene_list1 & result_record$CloneMutNum == 1 , "Clone_Independent & NoSMG_1Clone" , result_record$Filter2 )
## 所在克隆簇和已报道基因一起出现
result_record$Filter2 <- ifelse( !(result_record$gene %in% gene_list1) , "Clone_ConSMG" , result_record$Filter2 )
result_record$Filter2 <- ifelse( result_record$gene %in% dat_smg$Gene_Symbol , "SMG" , result_record$Filter2 )

##############################################################################
## 20230314
## 在其它同一类型肿瘤样本中存在至少1例突变CCF>0.6
result_record$AddtionalCloneSample <- ""
result_record$AddtionalCloneSampleNum <- ""

for( i in 1:nrow(result_record) ){

  geneN <- result_record$gene[i]
  normal_list <- unlist(strsplit(result_record$Normal[i] , ","))

  ## 判定该克隆簇所属肿瘤类型
  tumor_type <- result_record$Class[i]

  ## 所有该类型的样本,排除当前样本
  tumor_list <- subset( info , Class %in% tumor_type & !(Normal %in% normal_list) )

  ## 肿瘤中存在克隆选择的克隆簇
  tmp_ccf <- subset( dat_ccf , Hugo_Symbol == geneN & Variant_Classification %in% Variant_Type & CCF_adj >= clone_t )

  ## 判断样本的数量
  sample_name <- tmp_ccf[which(tmp_ccf$Sample %in% tumor_list$Tumor),"Sample"]
  sample_name <- unique(info[info$Tumor %in% sample_name,"ID"])

  result_record$AddtionalCloneSample[i] <- paste0(sample_name , collapse = ",")
  result_record$AddtionalCloneSampleNum[i] <- length(sample_name)
}

result_record$Filter2 <- ifelse( result_record$Filter2== "Clone_Independent & NoSMG_1Clone" & result_record$AddtionalCloneSampleNum >= 2 , 
  "Clone_Independent & NoSMG_1Clone & AddtionalCloneSample" , result_record$Filter2 )
result_record <- result_record[,c(1:4,7,8,5,6)]

images_name <- paste0(output_path , "/GCClone_gene.reord.tsv" )
write.table(result_record , images_name , sep = '\t' , quote=F , row.names=F)
